import torch
import json

stride = 2
padding = 0
kernalSize = 2
chanels = 3
imgSize =32

input = torch.randn(chanels, imgSize, imgSize)
operator = torch.nn.MaxPool2d(kernalSize,stride,padding)
output = operator.forward(input)

exportData = [
    ('input', input),
    ('output', output)
]

data = []
for name, tensor in exportData:
    data.append({"layer": name, "value": tensor.tolist()})

data.append({'layer': 'stride', 'value': stride})
data.append({'layer': 'padding', 'value': padding})
data.append({'layer': 'kernalSize', 'value': kernalSize})


jsonString = json.dumps(data)
with open('MaxPooling2DTest.json', 'w') as f:
    f.write(jsonString)
print('TestBench generate finished.')
